"""
Wan2.1-T2V-14B 推理引擎
使用 MindIE 进行文生视频推理
参考: https://www.hiascend.com/document/detail/zh/mindie/22RC1/quickstart/mindiesd_quickstart_0004.html
"""
import os
import glob
import logging
import subprocess
from typing import Dict, Optional
from datetime import datetime

logger = logging.getLogger(__name__)


class Wan2Inference:
    """Wan2.1-T2V推理引擎"""
    
    def __init__(self, model_loader, config: Dict = None):
        """
        初始化推理引擎
        
        Args:
            model_loader: Wan2ModelLoader 对象
            config: 配置字典
        """
        self.model_loader = model_loader
        self.config = config or {}
        
        # 获取配置
        api_config = self.config.get('api', {})
        self.output_dir = api_config.get('output_dir', '/app/outputs')
        self.temp_dir = api_config.get('temp_dir', '/app/temp')
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.temp_dir, exist_ok=True)
        
        # 获取推理参数
        self.inference_params = model_loader.get_inference_params()
    
    def generate_video(
        self,
        prompt: str,
        num_inference_steps: Optional[int] = None,
        seed: Optional[int] = None,
        request_id: Optional[str] = None,
        image_path: Optional[str] = None
    ) -> Dict:
        """
        文生视频/图生视频推理
        
        Args:
            prompt: 文本提示词
            num_inference_steps: 采样步数
            seed: 随机种子
            request_id: 请求ID
            image_path: 输入图片路径（可选，用于图生视频）
        
        Returns:
            包含视频路径的字典
        """
        try:
            # 生成请求ID
            if request_id is None:
                request_id = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            
            logger.info(f"开始视频生成任务 {request_id}")
            logger.info(f"提示词: {prompt}")
            if image_path:
                logger.info(f"输入图片: {image_path}")
            
            # 创建任务专属输出目录: /app/outputs/$request_id
            task_output_dir = os.path.join(self.output_dir, request_id)
            os.makedirs(task_output_dir, exist_ok=True)
            logger.info(f"创建输出目录: {task_output_dir}")
            
            # 准备推理参数
            steps = num_inference_steps or self.inference_params['sample_steps']
            base_seed = seed or 0
            
            # 构建torchrun命令
            cmd = self._build_inference_command(
                prompt=prompt,
                steps=steps,
                base_seed=base_seed,
                image_path=image_path
            )
            
            logger.info(f"执行推理命令: {' '.join(cmd)}")
            logger.info(f"工作目录: {task_output_dir}")
            
            # 在任务输出目录下执行推理
            result = subprocess.run(
                cmd,
                cwd=task_output_dir,  # 在任务目录下执行，输出文件会生成在这里
                capture_output=True,
                timeout=1800  # 30分钟超时
            )
            
            if result.returncode != 0:
                # 解码错误输出，忽略无法解码的字节
                stderr_text = result.stderr.decode('utf-8', errors='replace')
                logger.error(f"推理失败: {stderr_text}")
                raise RuntimeError(f"推理失败: {stderr_text}")
            
            # 查找生成的mp4文件
            mp4_files = glob.glob(os.path.join(task_output_dir, "*.mp4"))
            
            if not mp4_files:
                raise FileNotFoundError(f"未找到输出视频文件，目录: {task_output_dir}")
            
            # 取第一个mp4文件作为输出
            output_video_path = mp4_files[0]
            logger.info(f"视频生成完成: {output_video_path}")
            
            return {
                "request_id": request_id,
                "status": "success",
                "video_path": output_video_path,
                "video_dir": task_output_dir,
                "prompt": prompt,
                "steps": steps,
                "seed": base_seed
            }
            
        except Exception as e:
            logger.error(f"视频生成失败: {str(e)}", exc_info=True)
            raise
    
    def _build_inference_command(
        self,
        prompt: str,
        steps: int,
        base_seed: int,
        image_path: Optional[str] = None
    ) -> list:
        """构建推理命令（使用torchrun进行多卡推理）"""
        params = self.inference_params
        model_path = self.model_loader.get_model_path()
        wan2_code_path = self.model_loader.get_wan2_code_path()
        
        # 使用torchrun进行多卡推理
        cmd = [
            'torchrun',
            f'--nproc_per_node={params["nproc_per_node"]}',
            f'{wan2_code_path}/generate.py',
            '--task', params['task'],
            '--size', params['size'],
            '--ckpt_dir', model_path,
            '--sample_steps', str(steps),
            f'--ulysses_size', str(params['ulysses_size']),
            '--prompt', prompt,
        ]
        
        # 添加输入图片参数（图生视频）
        if image_path:
            cmd.extend(['--image', image_path])
        
        # 添加FSDP和并行选项
        if params.get('dit_fsdp'):
            cmd.append('--dit_fsdp')
        if params.get('t5_fsdp'):
            cmd.append('--t5_fsdp')
        if params.get('vae_parallel'):
            cmd.append('--vae_parallel')
        
        # 添加Attention Cache优化
        if params.get('use_attention_cache'):
            cmd.append('--use_attentioncache')
            cmd.extend(['--start_step', str(params['start_step'])])
            cmd.extend(['--attentioncache_interval', str(params['attention_cache_interval'])])
            cmd.extend(['--end_step', str(params['end_step'])])
        
        return cmd
